import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import math
import seaborn as sns
import matplotlib as mpl
from sklearn.svm import SVR
from sklearn.model_selection import train_test_split
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.metrics import mean_squared_error

mpl.use('TkAgg')

df = pd.read_excel("data/Q3O.xlsx")

y1_list = df['y1'].values
y2_list = df['y2'].values
y_list = df['y'].values

def NMSE(list1, list2):
    list1 = np.array(list1)
    list2 = np.array(list2)
    # 计算均方误差MSE
    mse = np.mean((list1 - list2) ** 2)
    # 计算根均方误差RMSE
    rmse = np.sqrt(mse)
    return rmse

print(NMSE(y1_list,y_list))
print(NMSE(y2_list,y_list))

y = df['y']
X = df.drop(['date','y'],axis=1)

X_train, X_test, y_train, y_test = train_test_split(X, y, train_size=0.8)

estimator = GradientBoostingRegressor()
estimator.fit(X_train, y_train)

y_predict = estimator.predict(X_test)
error = math.sqrt(mean_squared_error(y_test, y_predict))
print("NMSE：\n", error)

y3_list = estimator.predict(X)

# 画图
plt.figure(figsize=(8, 6), dpi=300)
plt.plot(df.index, y_list, label='Y')
plt.plot(df.index,  y1_list, label='ARIMA')
plt.plot(df.index, y2_list, label='LSTM')
plt.gca().xaxis.set_visible(False)
plt.legend()
plt.savefig("img/Q3_pic2.png")

plt.figure(figsize=(8, 6), dpi=300)
plt.plot(df.index, y_list, label='Y')
plt.plot(df.index, y3_list, label="ARIMA_LSTM")
plt.gca().xaxis.set_visible(False)
plt.legend()
plt.savefig("img/Q3_pic3.png")
# plt.show()

#
# inter = pd.read_excel("data/ans.xlsx", nrows=12)
# print(inter)
# X_inter = pd.concat([inter['LSTM'], inter['ARIMA']], axis=1)
# print(X_inter)
# X_inter.rename(columns={'LSTM':'y1', 'ARIMA': 'y2'},inplace=True)
# print(X_inter)
# y_inter = estimator.predict(X_inter)
# print(y_inter)
